# @File    : datavisualX.py
import sys
from typing import Union

import PyQt5.QtWidgets
import matplotlib.pyplot as plt
from PyQt5.Qt import QApplication

import autogen.data_visualization as data_visual
import common
import ezqt
import painter


class DataVisualX(data_visual.Ui_DataVisualization):
    PLOT_WHEN_RECEIVE_N_UPDATE_PLOT_MESSAGE = 2  # 每当收到这么多次更新绘图的消息后，才真正开始画一次。否则，可能会卡死。

    def __init__(self, hm: common.HallMachine,
                 parent: Union[PyQt5.QtWidgets.QWidget, ezqt.MainWindowKnowItselfShown],
                 timer: common.DataExchangerTimer = common.DEFAULT_MAIN_TIMER,
                 measurer: common.AbstractMeasurerTask = None, ):
        super(DataVisualX, self).__init__()
        self.__parent = parent
        super(DataVisualX, self).setupUi(self.__parent)
        self.hm = hm
        self.timer = timer
        self.measurer = measurer
        self.recorder = None
        self.canvas2d = ezqt.EzCanvasWidget(self.widget_for_2d_fig, plt.figure(1))
        self.canvas3d = ezqt.EzCanvasWidget(self.widget_for_3d_fig, plt.figure(2))

        self.d_components = {
            "Bx": common.RecorderBase.k_Bx_Gauss,
            "By": common.RecorderBase.k_By_Gauss,
            "Bz": common.RecorderBase.k_Bz_Gauss,
            "B": common.RecorderBase.k_B_Gauss,

            "Vx": common.RecorderBase.k_Vx_V,
            "Vy": common.RecorderBase.k_Vy_V,
            "Vz": common.RecorderBase.k_Vz_V,
            "V": common.RecorderBase.k_V_V,

            self.radioButton_t.text(): common.RecorderBase.k_timestamp,
            self.radioButton_s.text(): common.RecorderBase.k_s_mm,
            self.radioButton_X.text(): common.RecorderBase.k_x_mm,
            self.radioButton_Y.text(): common.RecorderBase.k_y_mm,
            self.radioButton_Z.text(): common.RecorderBase.k_z_mm,
        }
        self.__update_after_receive_n_message_from_now = DataVisualX.PLOT_WHEN_RECEIVE_N_UPDATE_PLOT_MESSAGE  # 再收到这么多次消息，就更新图
        self.connect_all()

    def connect_all(self):
        self.pushButton_plot_old.clicked.connect(lambda: self.__choose_old_record())
        self.timer.sig_plot_control_pts.connect(self.__on_plot_traj)
        self.pushButton_refresh.clicked.connect(self.__on_refresh)
        self.pushButton_int_along_graph_x.clicked.connect(self.__integrate_along_traj)
        self.pushButton_max_min_avg.clicked.connect(self.__max_min_avg)

    def __choose_old_record(self):
        filename, _ = PyQt5.Qt.QFileDialog.getOpenFileName(self.__parent, "选择一条之前的测量记录",
                                                           common.Recorder.DEFAULT_RECORD_DIR, "测量记录(*.csv)")
        if filename:
            temp_recorder = common.RecorderBase(lambda: filename)
            temp_recorder.read_csv()

            self.recorder = temp_recorder
            self.__on_refresh()

    def __on_refresh(self):
        if self.recorder:
            self.__plot_record(self.recorder, imediate=True)
            self.__integrate_along_traj()
            self.__max_min_avg()

    def __on_plot_traj(self, traj: common.Trajectory, ):
        """

        :param traj:
        :return:
        """
        painter.plot_traj_with_Traj(self.canvas3d.canvas_qtagg.figure.gca(projection="3d"), traj, True, "-",
                                    label="target traj")
        self.canvas3d.canvas_qtagg.draw()
        self.canvas3d.canvas_qtagg.flush_events()

    def set_measurer(self, measurer: common.AbstractMeasurerTask):
        self.measurer = measurer
        self.measurer.sig_mean_data_updated.connect(self.__plot_record)
        self.measurer.sig_raw_data_done.connect(self.__on_measurer_done)

    def __on_measurer_done(self, recorder_raw):
        self.__plot_record(recorder_raw, imediate=True)
        self.__on_refresh()

    def __get_x_column_of_2d_fig(self) -> str:
        """
        获取2D绘图区选中的横轴
        :return:
        """
        column_as_x = common.RecorderBase.k_timestamp
        for radio in self.widget_vars_as_x.children():
            if type(radio) == PyQt5.Qt.QRadioButton and radio.isChecked():
                column_as_x = self.d_components[radio.text()]
                break
        return column_as_x

    def __get_y_columns_of_2d_fig(self) -> set:
        """
        获取2D绘图区选中的纵轴
        :return:
        """
        components_to_draw = set()
        for checkbox in self.widget_vars_as_y.children():
            # checkbox:PyQt5.Qt.QCheckBox
            if type(checkbox) == PyQt5.QtWidgets.QCheckBox and checkbox.isChecked():
                components_to_draw.add(self.d_components[checkbox.text()])
        return components_to_draw

    def __plot_record(self, recorder: common.RecorderBase, imediate=False):
        """
        更新绘图区
        :param recorder:
        :return:
        """
        self.__update_after_receive_n_message_from_now -= 1
        if imediate or self.__update_after_receive_n_message_from_now <= 0:
            # 开始画图
            ax = self.canvas2d.canvas_qtagg.figure.gca()

            components_to_draw = self.__get_y_columns_of_2d_fig()
            column_as_x = self.__get_x_column_of_2d_fig()

            painter.plot_component_using_Record(ax, recorder, components_to_draw, column_as_x, clear_axes=True)
            self.canvas2d.canvas_qtagg.draw()
            self.canvas2d.canvas_qtagg.flush_events()

            self.__update_after_receive_n_message_from_now = DataVisualX.PLOT_WHEN_RECEIVE_N_UPDATE_PLOT_MESSAGE
            self.recorder = recorder

    def __integrate_along_traj(self):
        """
        沿着图中的横轴积分
        :return:
        """
        if self.recorder:
            column_as_x = self.__get_x_column_of_2d_fig()
            columns_as_y = list(self.__get_y_columns_of_2d_fig())

            recorder = self.recorder.integrate_value_along(columns_as_y, column_as_x)
            res = "沿%s积分结果：\n%s" % (column_as_x, recorder.df[columns_as_y])
            self.textEdit_integration.setText(res)

    def __max_min_avg(self):
        if self.recorder:
            columns_as_y = list(self.__get_y_columns_of_2d_fig())
            res = "最大值：\n%s\n最小值：\n%s\n平均值：\n%s" % (
                str(self.recorder.df.max()[columns_as_y]), self.recorder.df.min()[columns_as_y],
                self.recorder.df.mean()[columns_as_y])
            self.textEdit_max_min_avg.setText(res)


if __name__ == '__main__':
    app = QApplication([])

    main_window = PyQt5.QtWidgets.QMainWindow()
    widget = PyQt5.QtWidgets.QWidget(main_window)

    main_window.setCentralWidget(widget)

    hm = common.HallMachine()
    DataVisualX(hm, widget, )

    main_window.show()
    sys.exit(app.exec_())
